from collections import defaultdict
import os
import random
import shutil
from tqdm import tqdm

data_dir = '../../data/train'
filename = data_dir + '/label.txt'
save_path = '../../data/train_cls'
count_image = defaultdict(list)  # 所有的数据集 dict {'0':['1.png','2.png',...],'1':[],...}
with open(filename, 'r') as file_to_read:
    while True:
        lines = file_to_read.readline()
        if not lines:
            break
        (img_name, img_label) = lines.strip().split(':')
        count_image[img_label].append(img_name)
    ID = 0
    for idx, img_list in tqdm(count_image.items()):
        temp_save_path = save_path + '/' + str(ID)
        if not os.path.exists(temp_save_path):
            os.makedirs(temp_save_path)
        for image_name in img_list:
            query_path = os.path.join(data_dir+'/images', image_name)
            shutil.copy(query_path, temp_save_path)
        ID += 1
